package com.zhuang.treemap;

import java.util.Comparator;

public class TreeMap<K, V> {

	// 定义比较器变量
	private final Comparator<? super K> comparator;

	// 根结点
	private Entry<K, V> root;

	// 定义集合长度
	private int size;

	// 空参构造
	public TreeMap() {
		comparator = null;
	}

	// 有参构造
	public TreeMap(Comparator<? super K> comparator) {
		this.comparator = comparator;
	}

	// 定义内部类表示键值对
	private class Entry<K, V> {
		// 键
		K k;
		// 值
		V v;
		// 左子结点
		Entry<K, V> left;
		// 右子结点
		Entry<K, V> right;
		// 父结点
		Entry<K, V> parent;

		// 有参构造

		public Entry(K k, V v, Entry<K, V> left, Entry<K, V> right, Entry<K, V> parent) {
			this.k = k;
			this.v = v;
			this.left = left;
			this.right = right;
			this.parent = parent;
		}
	}

	// 获取集合长度
	public int size() {
		return size;
	}

	// get()方法的实现
	public V get(K key) {
		Entry<K, V> entry = getEntry(key);
		return entry == null ? null : entry.v;
	}

	// 根据键获取Entry对象的方法
	private Entry<K, V> getEntry(Object key) {
		// 非空校验
		if (key == null) {
			throw new NullPointerException();
		}

		// 给根结点起名
		Entry<K, V> t = root;

		// 判断有没有传入比较器
		if (comparator != null) {
			K k = (K) key;
			// 循环
			while (t != null) {
				int cmp = comparator.compare(k, t.k);
				if (cmp < 0) {
					t = t.left;
				} else if (cmp > 0) {
					t = t.right;
				} else {
					return t;
				}
			}

		} else {
			Comparable<? super K> k = (Comparable<? super K>) key;

			while (t != null) {
				int cmp = k.compareTo(t.k);
				if (cmp < 0) {
					t = t.left;
				} else if (cmp > 0) {
					t = t.right;
				} else {
					return t;
				}
			}
		}

		// 如果找不到返回null
		return null;

	}

	// put()方法的实现
	public V put(K key, V value) {
		// 给根结点赋值
		Entry<K, V> t = root;

		// 非空校验
		if (key == null) {
			throw new NullPointerException();
		}

		// 集合是否为空
		if (t == null) {
			// 创建新结点
			Entry<K, V> entry = new Entry<>(key, value, null, null, null);
			// 给根结点赋值
			root = entry;
			// 集合长度加一
			size++;
			return null;
		}

		// 键值对表示新增结点的父结点
		Entry<K, V> parent = t;
		// 定义变量
		int cmp = 0;
		if (comparator != null) {
			while (t != null) {
				// 给parent
				parent = t;
				// 判断键
				cmp = comparator.compare(key, t.k);

				if (cmp < 0) {
					t = t.left;
				} else if (cmp > 0) {
					t = t.right;
				} else {
					// 用新的值替换旧的值 把旧的值作为返回值返回
					V v = t.v;
					t.v = value;
					return v;
				}
			}
		} else {
			Comparable<? super K> k = (Comparable<? super K>) key;
			while (t != null) {
				parent = t;

				cmp = k.compareTo(t.k);

				if (cmp < 0) {
					t = t.left;
				} else if (cmp > 0) {
					t = t.right;
				} else {
					// 用新的值替换旧的值 把旧的值返回
					V v = t.v;
					t.v = value;
					return v;
				}
			}
		}

		// 要添加的键值对 键不重复
		Entry<K, V> entry = new Entry<>(key, value, null, null, parent);

		if (cmp < 0) {
			parent.left = entry;
		} else {
			parent.right = entry;
		}

		// 集合长度加一
		size++;

		return null;
	}

	// remove()方法的实现
	public V remove(K key) {

		Entry<K, V> entry = getEntry(key);

		if (entry == null) {
			return null;
		}

		// 删除操作
		// 有右子结点没有左子结点
		if (entry.left == null && entry.right != null) {
			if (entry == root) {
				// 删除节点的右子节点直接成为根节点
				root = entry.right;
			} else if (entry.parent.right == entry) {
				// 原来节点的右子节点赋值给父节点的右子节点
				entry.parent.right = entry.right;
			} else if (entry.parent.left == entry) {
				// 原来节点的右子节点赋值给父节点的左子节点
				entry.parent.left = entry.right;
			}

			// 让被删除结点的子结点指向父结点
			entry.right.parent = entry.parent;
			// 有左子结点没有右子结点
		} else if (entry.left != null && entry.right == null) {
			if (entry == root) {
				// 删除节点的左子节点直接成为根节点
				root = entry.left;
			} else if (entry.parent.right == entry) {
				// 原来节点的左子节点赋值给父节点的右子节点
				entry.parent.right = entry.left;
			} else if (entry.parent.left == entry) {
				// 原来节点的左子节点赋值给父节点的左子节点
				entry.parent.left = entry.left;
			}

			// 让被删除结点子结点指向父结点
			entry.left.parent = entry.parent;

			// 有左子结点也有右子结点
		} else if (entry.left != null && entry.right != null) {

			// 要后继结点
			Entry<K, V> target = entry.right;
			// 寻找被删除结点右子结点最左子结点
			while (target.left != null) {
				target = target.left;
			}

			// 右子结点作为后继结点
			if (entry.right == target) {
				// 下边节点直接指向上边节点
				target.parent = entry.parent;

				if (entry == root) {
					root = target;
					// 判断当前删除节点是父节点的右节点
				} else if (entry.parent.right == entry) {
					// 父节点的右节点指向target
					entry.parent.right = target;
					// 判断当前删除节点是父节点的左节点
				} else if (entry.parent.left == entry) {
					// 父节点的左节点指向target
					entry.parent.left = target;
				}
				// 被删除结点左子结点重新指向新的父结点
				entry.left.parent = target;
				target.left = entry.left;

			} else {
				// 右子结点的最左子结点作为后继结点
				if (target.right == null) {
					// 后继结点没有右子结点
					target.parent.left = null;
				} else {
					// 后继结点有右子结点
					target.parent.left = target.right;
					target.right = target.parent;
				}

				// 让后继结点替换被删除结点
				if (entry == root) {
					root = target;
				} else if (entry.parent.right == entry) {
					entry.parent.right = target;
				} else if (entry.parent.left == entry) {
					entry.parent.left = target;
				}

				// 被删除结点左右子树需要指向后继结点
				entry.left.parent = target;
				entry.right.parent = target;
				target.left = entry.left;
				target.right = entry.right;
			}

		} else {
			// 要删除的结点是叶子结点
			if (entry == root) {
				root = null;
			} else if (entry.parent.right == entry) {
				entry.parent.right = null;
			} else if (entry.parent.left == entry) {
				entry.parent.left = null;
			}
		}

		// 给集合长度减一
		size--;

		return entry.v;
	}

	// toString()方法的实现
	// {1=abc, 3=qwe}
	public String toString() {
		// 非空判断
		if (root == null) {
			return "{}";
		}

		String s = "{";
		String s1 = method(root);

		s = s + s1.substring(0, s1.length() - 2) + "}";

		return s;
	}

	// 递归的方法
	// 1=abc, 3=qwe,
	private String method(Entry<K, V> entry) {

		String s = "";

		// 拼接左子结点
		if (entry.left != null) {
			s += method(entry.left);
		}

		// 拼接中间结点自己
		s += entry.k + "=" + entry.v + ", ";

		// 拼接右子结点
		if (entry.right != null) {
			s += method(entry.right);
		}

		return s;
	}

}
